package server

import "C"
import (
	"bytes"
	"encoding/binary"
	"net"
	"syscall"
	"unsafe"

	"github.com/cilium/ebpf"
	"github.com/gobuffalo/packr"
	"github.com/kproxy/kproxy.go/pkg/logging"
	"github.com/kproxy/kproxy.go/pkg/utils"
	"github.com/pkg/errors"
)

/*
#include <linux/types.h>
typedef struct sock_key {
    __u32 remote_ip;
    __u32 remote_port;
    __u32 local_ip;
    __u32 local_port;
} sock_key;
*/
import "C"

var logger = logging.RegisterScope("server")

const sockopsDir = "../sockops/"

var box = packr.NewBox(sockopsDir)

type ServerArgs struct {
	Address string
}

type Server struct {
	args       *ServerArgs
	listener   net.Listener
	epoller    *Epoll
	connset    *ConnSet
	waitgroup  *utils.WaitGroup
	collection *ebpf.Collection
}

func NewServer(args *ServerArgs) (*Server, error) {
	listener, err := net.Listen("tcp", args.Address)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	epoller, err := NewEpoll()
	if err != nil {
		return nil, err
	}
	server := &Server{
		args:      args,
		listener:  listener,
		epoller:   epoller,
		connset:   NewConnSet(),
		waitgroup: utils.NewWaitGroup(),
	}
	redirectFileData, err := box.Find("bpf_redirect.o")
	if err != nil {
		return nil, err
	}
	collectionSpec, err := ebpf.LoadCollectionSpecFromReader(bytes.NewReader(redirectFileData))
	if err != nil {
		return nil, err
	}
	if server.collection, err = ebpf.NewCollection(collectionSpec); err != nil {
		return nil, err
	}
	if err := server.collection.Programs["prog_parser"].Attach(server.collection.Maps["sock_map"].FD(), ebpf.AttachSkSKBStreamParser, 0); err != nil {
		return nil, err
	}
	if err := server.collection.Programs["prog_verdict"].Attach(server.collection.Maps["sock_map"].FD(), ebpf.AttachSkSKBStreamVerdict, 0); err != nil {
		return nil, err
	}

	return server, nil
}

func (server *Server) acceptLoop() {
	logger.Infof("listening on %s\n", server.args.Address)
	for {
		conn, err := server.listener.Accept()
		if err != nil {
			conn.Close()
			continue
		}
		if err := server.connset.Add(conn); err != nil {
			conn.Close()
			continue
		}
		//key := 0
		remoteTcpAddr := conn.RemoteAddr().(*net.TCPAddr)
		localTcpAddr := conn.LocalAddr().(*net.TCPAddr)

		remoteIP := binary.BigEndian.Uint32(reverseBytes(remoteTcpAddr.IP.To4()))
		b := make([]byte, 4)
		binary.LittleEndian.PutUint32(b, uint32(remoteTcpAddr.Port))
		remotePort := binary.BigEndian.Uint32(b)

		localIP := binary.BigEndian.Uint32(reverseBytes(localTcpAddr.IP.To4()))
		localPort := localTcpAddr.Port

		var key C.sock_key

		key.remote_ip = C.__u32(remoteIP)
		key.remote_port = C.__u32(remotePort)
		key.local_ip = C.__u32(localIP)
		key.local_port = C.__u32(localPort)
		file, err := conn.(*net.TCPConn).File()
		if err != nil {
			logger.Error(errors.WithStack(err))
			continue
		}

		socketFd := file.Fd()
		server.collection.Maps["sock_map"].Update(key, unsafe.Pointer(&socketFd), ebpf.UpdateAny)
		if err := server.epoller.Add(conn); err != nil {
			server.epoller.Remove(conn)
			conn.Close()
			continue
		}
		// TODO add sock_hash_update_elem
	}
}
func (server *Server) processConns() {
	for {
		events, err := server.epoller.Wait()
		if err != nil && err != syscall.EINTR {
			logger.Errorf("failed to poll event: %v, events len: %d", err, len(events))
			continue
		}

		for _, event := range events {
			if int(event.Events)&syscall.EPOLLIN > 0 {
				conn := server.epoller.GetConn(int(event.Fd))
				if conn == nil {
					continue
				}
				b := make([]byte, 1024)
				n, err := conn.Read(b)
				if err == nil {
					if _, err := conn.Write(b); err != nil {
						logger.Errorf("write error %v\n", err)
					}
				} else if n == 0 || err != nil {
					if err == syscall.EAGAIN {
						continue
					}
					if err := server.epoller.Remove(conn); err != nil {
						logger.Errorf("failed to remove conn: %s \n", err)
						continue
					}
					conn.Close()
				} else {
					logger.Errorf("read err: ", err)
				}
			}
		}
	}
}

func (server *Server) Run() error {
	server.waitgroup.GoFunc(server.acceptLoop)
	server.waitgroup.GoFunc(server.processConns)
	server.waitgroup.Wait()
	return nil
}

func reverseBytes(src []byte) []byte {
	dst := make([]byte, len(src))
	copy(dst, src)
	for i := 0; i < len(dst)/2; i++ {
		j := len(dst) - i - 1
		dst[i], dst[j] = dst[j], dst[i]
	}
	return dst
}
